# sfb
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import itertools
from sklearn.metrics import r2_score
from matplotlib.pylab import style  # 自定义图表风格
import statsmodels.api as sm  # D-W检验,一阶自相关检验

style.use('ggplot')

# 解决中文乱码问题
plt.rcParams['font.sans-serif'] = ['Simhei']
# 解决坐标轴刻度负号乱码
plt.rcParams['axes.unicode_minus'] = False


# @param: ratio - 比率 data - 数据, @return: predictResultData
def predictSale(data, ratio):
    print("------------- predictSale method is starting -------------")

    """ 获取数据 start """
    month = data["yearAndMonth"]
    sale = data["salesData"]
    seriesName = data["series"]
    seriesId = data["seriesId"]
    # 计算训练数据大小
    train_size = round(len(month) * ratio)

    # 判断0的个数，多于十二个0的数据不做计算
    if sum(i == '0' for i in sale) > 12:
        return
    """ 获取数据 end """

    # 数据预处理
    p_sale, train, test = preHandleData(month, sale, train_size)
    # 学习参数
    param = getParam(p_sale)
    # 滚动预测 + 传统预测
    y_true, y_pred = rollPredict(train, test, param)
    y_tradition = traditionPredict(train, param, len(sale) - train_size)

    # 进行R平方打分
    g_score = r2_score(np.array(y_true), np.array(y_pred))
    t_score = r2_score(np.array(y_true), y_tradition.values)
    print("滚动预测R平方相关系数：%s,\n传统预测R平方相关系数：%s" % (g_score, t_score))

    print("------------- predictSale method is ending   -------------")
    return {
        "seriesId": seriesId,
        "seriesName": seriesName,
        "salesData": sale,
        "rollPredRes": [round(x, 2) for x in y_pred],
        "rollScore": round(g_score, 2),
        "traditionPredRes": [round(x, 2) for x in y_tradition.values.tolist()],
        "traditionScore": round(t_score, 2),
        "pdq": param.values[0].tolist()[0],
        "pdq_x_PDQs": param.values[0].tolist()[1],
        "aic": round(param.values[0].tolist()[2], 2),
    }


# @param: forcastMonth - 比率 data - 数据, @return: predictResultData
def predictSaleTradition(data, forcastMonth):
    print("------------- predictSaleTradition method is starting -------------")

    """ 获取数据 start """
    month = data["yearAndMonth"]
    sale = data["salesData"]
    seriesName = data["series"]
    seriesId = data["seriesId"]

    # 判断0的个数，多于十二个0的数据不做计算
    if sum(i == '0' for i in sale) > 12:
        return
    """ 获取数据 end """

    # 数据预处理
    p_sale = preHandleDataTradition(month, sale)
    # 学习参数
    param = getParam(p_sale)
    # 滚动预测 + 传统预测
    y_tradition = traditionPredict(p_sale, param, forcastMonth)

    print("------------- predictSale method is ending   -------------")
    return {
        "seriesId": seriesId,
        "seriesName": seriesName,
        "salesData": sale,
        "traditionPredRes": [round(x, 2) for x in y_tradition.values.tolist()],
        "pdq": param.values[0].tolist()[0],
        "pdq_x_PDQs": param.values[0].tolist()[1],
        "aic": round(param.values[0].tolist()[2], 2),
    }


# @Introduce: 数据预处理
# @Param: month - 月份, sale - 销量, ratio - 训练集比率, train_size - 训练集大小
# @Return p_sale - 完整销量数据, train - 训练集, test - 测试集
def preHandleData(month, sale, train_size):
    # 创建dataframe
    p_sale = pd.DataFrame({"date": month, "sale": sale})
    # 转换日期数据类型 和 销售量数据类型
    p_sale['date'] = pd.to_datetime(p_sale['date'])
    p_sale.sale = p_sale.sale.astype('float')
    # 设置索引值
    p_sale.set_index('date', inplace=True)

    # 获取训练、测试数据集
    train = p_sale[:train_size]
    test = p_sale[train_size:]

    return p_sale, train, test


# @Introduce: 数据预处理传统预测
# @Param: month - 月份, sale - 销量
# @Return p_sale - 完整销量数据
def preHandleDataTradition(month, sale):
    # 创建dataframe
    p_sale = pd.DataFrame({"date": month, "sale": sale})
    # 转换日期数据类型 和 销售量数据类型
    p_sale['date'] = pd.to_datetime(p_sale['date'])
    p_sale.sale = p_sale.sale.astype('float')
    # 设置索引值
    p_sale.set_index('date', inplace=True)

    return p_sale


# @Introduce: 学习参数
# @Param: p_sale - 销量
# @Return param - aic最小的组合
def getParam(p_sale):
    # 设定pdq大小在 [0,2) 之间
    p = d = q = range(0, 2)

    # 生成普通 p,d,q 的组合
    pdq = list(itertools.product(p, d, q))

    # 生成季节性 P,D,Q,s 的组合
    pdq_x_PDQs = [(x[0], x[1], x[2], 12) for x in list(itertools.product(p, d, q))]

    # 自动学习出aic最小的参数组合
    a = []
    b = []
    c = []
    wf = pd.DataFrame()
    for param in pdq:
        for seasonal_param in pdq_x_PDQs:
            try:
                mod = sm.tsa.statespace.SARIMAX(p_sale, order=param, seasonal_order=seasonal_param,
                                                enforce_stationarity=False, enforce_invertibility=False)
                results = mod.fit()
                # print('ARIMA{}x{} - AIC:{}'.format(param, seasonal_param, results.aic))
                a.append(param)
                b.append(seasonal_param)
                c.append(results.aic)
            except:
                continue
    wf['pdq'] = a
    wf['pdq_x_PDQs'] = b
    wf['aic'] = c

    # 保存 aic 最小的参数组合
    param = wf[wf['aic'] == wf['aic'].min()]
    return param


# @Introduce: 滚动预测
# @Param: train - 训练集, test - 测试集, param - 参数组合
# @Return y_true - 真实值, y_pred - 预测结果
def rollPredict(train, test, param):
    # 拿到训练数据集的销量
    history_endog = list(train['sale'].copy(deep=True))

    # 真实销量值
    y_true = []
    # 预测销量值
    y_pred = []

    # 每次根据历史销量预测下一个月的销量，每次加入一个月真实测试值，实现滚动预测
    for obs in test['sale']:
        # 训练模型
        temp_model = sm.tsa.statespace.SARIMAX(history_endog,
                                               order=param['pdq'].values[0],
                                               seasonal_order=param['pdq_x_PDQs'].values[0])

        model_fit = temp_model.fit()
        # 将测试集中真实销量加入真实集中
        y_true.append(obs)
        # 将下一月的预测值加入预测集中
        y_pred.append(model_fit.forecast()[0])
        # 往训练集中加入下一月的真实销量
        history_endog.append(obs)

    return y_true, y_pred


# @Introduce: 传统预测
# @Param: train - 训练集, param - 参数组合, size - 需要预测大小
# @Return y_tradition - 预测结果
def traditionPredict(train, param, size):
    model = sm.tsa.statespace.SARIMAX(train['sale'],
                                      order=param['pdq'].values[0],
                                      seasonal_order=param['pdq_x_PDQs'].values[0])

    rs = model.fit()
    y_tradition = rs.forecast(size)
    return y_tradition
